package spring.my.db.sharding.demo.config.db;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.springframework.stereotype.Component;
import org.springframework.util.ReflectionUtils;

import javax.sql.DataSource;
import java.lang.reflect.Field;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;


public class DynamicDataSource extends AbstractRoutingDataSource {
    private final Logger log = LoggerFactory.getLogger(this.getClass());

    private Integer slaveCount = 0;

    private AtomicInteger counter = new AtomicInteger(0);

    private List<Object> slaveDataSources = new ArrayList<>();

    @Override
    protected Object determineCurrentLookupKey() {
        if (DynamicDataSourceHolder.isMaster) {
            return DynamicDataSourceHolder.getDataSource();
        }
        Object slaveKey = getSlaveKey();
        if (slaveKey != null) {
            log.info(Thread.currentThread().getName() + " [slave-key = {}]", slaveKey);
            return slaveKey;
        }
        return DynamicDataSourceHolder.getDataSource();
    }

    @Override
    public void afterPropertiesSet() {
        super.afterPropertiesSet();
        // 获取 AbstractRoutingDataSource 中的属性
        Field field = ReflectionUtils.findField(AbstractRoutingDataSource.class, "resolvedDataSources");
        if (field != null) {
            // 设置属性可见
            field.setAccessible(true);
            try {
                Map<String, Object> resolvedDataSources = (Map<String, Object>) field.get(this);
                // 记录主数据库之后的从数据库
                this.slaveCount = resolvedDataSources.size() - 1;
                resolvedDataSources.forEach((k, v) -> {
                    if (!DynamicDataSourceHolder.MASTER.equals(k)) {
                        log.info("slave = [key = {}]", k);
                        slaveDataSources.add(k);
                    }
                });
            } catch (IllegalAccessException e) {
                log.error("afterPropertiesSet error.");
                e.printStackTrace();
            }
        }
    }

    public Object getSlaveKey() {
        if (this.slaveCount > 0) {
            int index = counter.getAndIncrement() % this.slaveCount;
            if (index >= this.slaveCount) {
                counter.set(0);
            }
            return slaveDataSources.get(index);
        }
        return null;
    }
}

